{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SageMaker Benchmark Evaluation - Basic Usage\n",
    "\n",
    "This notebook demonstrates the basic user-facing flow for creating and managing benchmark evaluation jobs using the BenchmarkEvaluator with Jinja2 template-based pipeline generation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Discover Available Benchmarks\n",
    "\n",
    "Discover the benchmark properties and available options:\n",
    "https://docs.aws.amazon.com/sagemaker/latest/dg/nova-model-evaluation.html"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Configure AWS credentials and region\n",
    "#! ada credentials update --provider=isengard --account=<> --role=Admin --profile=default --once\n",
    "#! aws configure set region us-west-2"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from sagemaker.train.evaluate import get_benchmarks, get_benchmark_properties\n",
    "from rich.pretty import pprint\n",
    "\n",
    "# Configure logging to show INFO messages\n",
    "import logging\n",
    "logging.basicConfig(\n",
    "    level=logging.INFO,\n",
    "    format='%(levelname)s - %(name)s - %(message)s'\n",
    ")\n",
    "\n",
    "# Get available benchmarks\n",
    "Benchmark = get_benchmarks()\n",
    "pprint(list(Benchmark))\n",
    "\n",
    "# Print properties for a specific benchmark\n",
    "pprint(get_benchmark_properties(benchmark=Benchmark.GEN_QA))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Create BenchmarkEvaluator\n",
    "\n",
    "Create a BenchmarkEvaluator instance with the desired benchmark. The evaluator will use Jinja2 templates to render a complete pipeline definition.\n",
    "\n",
    "**Required Parameters:**\n",
    "- `benchmark`: Benchmark type from the Benchmark enum\n",
    "- `base_model`: Model ARN from SageMaker hub content\n",
    "- `output_s3_location`: S3 location for evaluation outputs\n",
    "- `mlflow_resource_arn`: MLflow tracking server ARN for experiment tracking\n",
    "\n",
    "**Optional Template Fields:**\n",
    "These fields are used for template rendering. If not provided, defaults will be used:\n",
    "- `model_package_group`: Model package group ARN\n",
    "- `source_model_package`: Source model package ARN\n",
    "- `dataset`: S3 URI of evaluation dataset\n",
    "- `model_artifact`: ARN of model artifact for lineage tracking (auto-inferred from source_model_package)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from sagemaker.train.evaluate import BenchMarkEvaluator\n",
    "\n",
    "# Create evaluator with GEN_QA benchmark\n",
    "# These values match our successfully tested configuration\n",
    "evaluator = BenchMarkEvaluator(\n",
    "    benchmark=Benchmark.GEN_QA,\n",
    "    model=\"arn:aws:sagemaker:us-west-2:<>:model-package/test-finetuned-models-gamma/28\",\n",
    "    s3_output_path=\"s3://mufi-test-serverless-smtj/eval/\",\n",
    "    mlflow_resource_arn=\"arn:aws:sagemaker:us-west-2:<>:mlflow-tracking-server/mmlu-eval-experiment\",\n",
    "    dataset=\"s3://sagemaker-us-west-2-<>/studio-users/d20251107t195443/datasets/2025-11-07T19-55-37-609Z/zc_test.jsonl\",\n",
    "    model_package_group=\"arn:aws:sagemaker:us-west-2:<>:model-package-group/example-name-aovqo\", # Optional inferred from model if model package\n",
    "    base_eval_name=\"gen-qa-eval-demo\",\n",
    "    # Note: sagemaker_session is optional and will be auto-created if not provided\n",
    "    # Note: region is optional and will be auto deduced using environment variables - SAGEMAKER_REGION, AWS_REGION\n",
    ")\n",
    "\n",
    "pprint(evaluator)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# # [Optional] BASE MODEL EVAL\n",
    "\n",
    "# from sagemaker.train.evaluate import BenchMarkEvaluator\n",
    "\n",
    "# # Create evaluator with GEN_QA benchmark\n",
    "# # These values match our successfully tested configuration\n",
    "# evaluator = BenchMarkEvaluator(\n",
    "#     benchmark=Benchmark.GEN_QA,\n",
    "#     model=\"meta-textgeneration-llama-3-2-1b-instruct\",\n",
    "#     s3_output_path=\"s3://mufi-test-serverless-smtj/eval/\",\n",
    "#     mlflow_resource_arn=\"arn:aws:sagemaker:us-west-2:<>:mlflow-tracking-server/mmlu-eval-experiment\",\n",
    "#     dataset=\"s3://sagemaker-us-west-2-<>/studio-users/d20251107t195443/datasets/2025-11-07T19-55-37-609Z/zc_test.jsonl\",\n",
    "#     # model_package_group=\"arn:aws:sagemaker:us-west-2:<>:model-package-group/example-name-aovqo\", # Optional inferred from model if model package\n",
    "#     base_eval_name=\"gen-qa-eval-demo\",\n",
    "#     # Note: sagemaker_session is optional and will be auto-created if not provided\n",
    "#     # Note: region is optional and will be auto deduced using environment variables - SAGEMAKER_REGION, AWS_REGION\n",
    "# )\n",
    "\n",
    "# pprint(evaluator)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# # [Optional] Nova testing IAD Prod\n",
    "\n",
    "# from sagemaker.train.evaluate import BenchMarkEvaluator\n",
    "\n",
    "# # Create evaluator with GEN_QA benchmark\n",
    "# # These values match our successfully tested configuration\n",
    "# evaluator = BenchMarkEvaluator(\n",
    "#     benchmark=Benchmark.GEN_QA,\n",
    "#     # model=\"arn:aws:sagemaker:us-east-1:<>:model-package/bgrv-nova-micro-sft-lora/1\",\n",
    "#     model=\"arn:aws:sagemaker:us-east-1:<>:model-package/test-nova-finetuned-models/3\",\n",
    "#     s3_output_path=\"s3://mufi-test-serverless-iad/eval/\",\n",
    "#     mlflow_resource_arn=\"arn:aws:sagemaker:us-east-1:<>:mlflow-tracking-server/mlflow-prod-server\",\n",
    "#     dataset=\"s3://sagemaker-us-east-1-<>/studio-users/d20251107t195443/datasets/2025-11-07T19-55-37-609Z/zc_test.jsonl\",\n",
    "#     model_package_group=\"arn:aws:sagemaker:us-east-1:<>:model-package-group/test-nova-finetuned-models\", # Optional inferred from model if model package\n",
    "#     base_eval_name=\"gen-qa-eval-demo\",\n",
    "#     region=\"us-east-1\",\n",
    "#     # Note: sagemaker_session is optional and will be auto-created if not provided\n",
    "#     # Note: region is optional and will be auto deduced using environment variables - SAGEMAKER_REGION, AWS_REGION\n",
    "# )\n",
    "\n",
    "# pprint(evaluator)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optionally update the hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "pprint(evaluator.hyperparameters.to_dict())\n",
    "\n",
    "# optionally update hyperparameters\n",
    "# evaluator.hyperparameters.temperature = \"0.1\"\n",
    "\n",
    "# optionally get more info on types, limits, defaults.\n",
    "# evaluator.hyperparameters.get_info()\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Run Evaluation\n",
    "\n",
    "Start a benchmark evaluation job. The system will:\n",
    "1. Build template context with all required parameters\n",
    "2. Render the pipeline definition from `DETERMINISTIC_TEMPLATE` using Jinja2\n",
    "3. Create or update the pipeline with the rendered definition\n",
    "4. Start the pipeline execution with empty parameters (all values pre-substituted)\n",
    "\n",
    "**What happens during execution:**\n",
    "- CreateEvaluationAction: Sets up lineage tracking\n",
    "- EvaluateBaseModel & EvaluateCustomModel: Run in parallel as serverless training jobs\n",
    "- AssociateLineage: Links evaluation results to lineage tracking"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Run evaluation with configured parameters\n",
    "execution = evaluator.evaluate()\n",
    "pprint(execution)\n",
    "\n",
    "print(f\"\\nPipeline Execution ARN: {execution.arn}\")\n",
    "print(f\"Initial Status: {execution.status.overall_status}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Alternative: Override Subtasks at Runtime\n",
    "\n",
    "For benchmarks with subtask support, you can override subtasks when calling evaluate():"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Override subtasks at evaluation time\n",
    "# execution = mmlu_evaluator.evaluate(subtask=\"abstract_algebra\")  # Single subtask\n",
    "# execution = mmlu_evaluator.evaluate(subtask=[\"abstract_algebra\", \"anatomy\"])  # Multiple subtasks"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Monitor Execution\n",
    "\n",
    "Check the job status and refresh as needed:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Refresh status\n",
    "execution.refresh()\n",
    "\n",
    "# Display job status with step details\n",
    "pprint(execution.status)\n",
    "\n",
    "# Display individual step statuses\n",
    "if execution.status.step_details:\n",
    "    print(\"\\nStep Details:\")\n",
    "    for step in execution.status.step_details:\n",
    "        print(f\"  {step.name}: {step.status}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Wait for Completion\n",
    "\n",
    "Wait for the pipeline to complete. This provides rich progress updates in Jupyter notebooks:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Wait for job completion with progress updates\n",
    "# This will show a rich progress display in Jupyter\n",
    "execution.wait(target_status=\"Succeeded\", poll=5, timeout=3600)\n",
    "\n",
    "print(f\"\\nFinal Status: {execution.status.overall_status}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: View Results\n",
    "\n",
    "Display the evaluation results in a formatted table:\n",
    "\n",
    "Output Structure:\n",
    "\n",
    "Evaluation results are stored in S3:\n",
    "\n",
    "```\n",
    "s3://your-bucket/output/\n",
    "└── job_name/\n",
    "    └── output/\n",
    "        └── output.tar.gz\n",
    "```\n",
    "\n",
    "Extract output.tar.gz to reveal:\n",
    "\n",
    "```\n",
    "run_name/\n",
    "├── eval_results/\n",
    "│   ├── results_[timestamp].json\n",
    "│   ├── inference_output.jsonl (for gen_qa)\n",
    "│   └── details/\n",
    "│       └── model/\n",
    "│           └── <execution-date-time>/\n",
    "│               └── details_<task_name>_#_<datetime>.parquet\n",
    "└── tensorboard_results/\n",
    "    └── eval/\n",
    "        └── events.out.tfevents.[timestamp]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "pprint(execution.s3_output_path)\n",
    "# Display results in a formatted table\n",
    "execution.show_results()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: Retrieve an Existing Job\n",
    "\n",
    "You can retrieve and inspect any existing evaluation job:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from sagemaker.train.evaluate import EvaluationPipelineExecution\n",
    "from rich.pretty import pprint\n",
    "\n",
    "\n",
    "# Get an existing job by ARN\n",
    "# Replace with your actual pipeline execution ARN\n",
    "existing_arn = \"arn:aws:sagemaker:us-west-2:<>:pipeline/SagemakerEvaluation-BenchmarkEvaluation-c344c91d-6f62-4907-85cc-7e6b29171c42/execution/inlsexrd7jes\"\n",
    "\n",
    "# base model only example\n",
    "# existing_arn = \"arn:aws:sagemaker:us-west-2:<>:pipeline/SagemakerEvaluation-benchmark/execution/gdp9f4dbv2vi\"\n",
    "existing_execution = EvaluationPipelineExecution.get(\n",
    "    arn=existing_arn,\n",
    "    region=\"us-west-2\"\n",
    ")\n",
    "\n",
    "pprint(existing_execution)\n",
    "print(f\"\\nStatus: {existing_execution.status.overall_status}\")\n",
    "\n",
    "existing_execution.show_results()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Run evaluation with configured parameters\n",
    "execution = evaluator.evaluate()\n",
    "pprint(execution)\n",
    "\n",
    "print(f\"\\nPipeline Execution ARN: {execution.arn}\")\n",
    "print(f\"Initial Status: {execution.status.overall_status}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 8: List All Benchmark Evaluations\n",
    "\n",
    "Retrieve all benchmark evaluation executions:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Get all benchmark evaluations (returns iterator)\n",
    "all_executions_iter = BenchMarkEvaluator.get_all(region=\"us-west-2\")\n",
    "all_executions = list(all_executions_iter)\n",
    "\n",
    "print(f\"Found {len(all_executions)} evaluation(s)\\n\")\n",
    "for exec in all_executions[:5]:  # Show first 5\n",
    "    print(f\"  {exec.name}: {exec.status.overall_status}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 9: Stop a Running Job (Optional)\n",
    "\n",
    "You can stop a running evaluation if needed:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Uncomment to stop the job\n",
    "# existing_execution.stop()\n",
    "# print(f\"Execution stopped. Status: {execution.status.overall_status}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding the Pipeline Structure\n",
    "\n",
    "The rendered pipeline definition includes:\n",
    "\n",
    "**4 Steps:**\n",
    "1. **CreateEvaluationAction** (Lineage): Sets up tracking\n",
    "2. **EvaluateBaseModel** (Training): Evaluates base model\n",
    "3. **EvaluateCustomModel** (Training): Evaluates custom model\n",
    "4. **AssociateLineage** (Lineage): Links results\n",
    "\n",
    "**Key Features:**\n",
    "- Template-based: Uses Jinja2 for flexible pipeline generation\n",
    "- Parallel execution: Base and custom models evaluated simultaneously\n",
    "- Serverless: No need to manage compute resources\n",
    "- MLflow integration: Automatic experiment tracking\n",
    "- Lineage tracking: Full traceability of evaluation artifacts\n",
    "\n",
    "**Typical Execution Time:**\n",
    "- Total: ~10-12 minutes\n",
    "- Downloading phase: ~5-7 minutes (model and dataset)\n",
    "- Training phase: ~3-5 minutes (running evaluation)\n",
    "- Lineage steps: ~2-4 seconds each"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
